# #!/bin/bash

gpu=0
dataset_dir='./data'

# ## This super-network will be generated after executing 1_script_train_supernet.sh
search_mode=random ## WS + random search (in Table 5)
# search_mode=evolution ## AutoSNN


searchString="DVS"
for dataset_name in CIFAR10 CIFAR100 CIFAR10DVS DVS128Gesture Tiny-ImageNet-200
do

isDVS=$(echo $dataset_name | grep "${searchString}")

if [[ "$isDVS" != "" ]] 
then
    echo "$dataset_name is DVS dataset"
    T=20
    channels=16
    suffix_prefix='_Adam_2022_T_'$T'_init_tau_2.0_vth_1.0_neuron_PLIF_split_by_number_normalization_None'
    batch_size=48
else
    echo "$dataset_name is not DVS dataset"
    T=8 #20
    channels=64 #16
    suffix_prefix='_SNN_Adam_600ep_2022'
    batch_size=96 #64
fi

suffix=$suffix_prefix'/batch_size_'$batch_size'/checkpoint.pth.tar'
search_space=AutoSNN_$channels
prefix='macro_search_result2/uniform_sampling/'$search_space'_'
trained_supernet=$prefix$dataset_name$suffix

python search_arch/search.py \
    --gpu $gpu \
    --T $T --init_tau 2.0 --v_threshold 1.0 --neuron PLIF \
    --dataset_dir $dataset_dir \
    --dataset_name $dataset_name \
    --supernet $trained_supernet \
    --seed 2022 \
    --search_space $search_space \
    --search_algo $search_mode \
    --fitness ACC_pow_spikes \
    --fitness_lambda -0.08 \
    --batch_size $batch_size

echo "Searching AutoSNN on $dataset_name is completed \n"

done